class Role(db.Document, RoleMixin):
    name = db.StringField(max_length=80, unique=True)
    description = db.StringField(max_length=255)

class User(db.Document, UserMixin):
    email = db.StringField(max_length=255)
    password = db.StringField(max_length=255)
    active = db.BooleanField(default=True)
    confirmed_at = db.DateTimeField()
    roles = db.ListField(db.ReferenceField(Role), default=[])

class MongoDatastore(Datastore):
    def put(self, model):
        model.save()
        return model

    def delete(self, model):
        model.delete()

class MongoUserDatastore(MongoDatastore, UserDatastore):
    """A Mongo datastore implementation for Flask-Security that assumes
    the use of the Flask-Mongo extension.
    """
    def __init__(self, db, user_model, role_model):
        MongoDatastore.__init__(self, db)
        UserDatastore.__init__(self, user_model, role_model)

    def get_user(self, id_or_email):
        from mongoengine import ValidationError
        try:
            return self.user_model.objects(id=id_or_email).first()
        except ValidationError:
            return self.user_model.objects(email__iexact=id_or_email).first()

    def find_user(self, **kwargs):
        try:
            from mongoengine.queryset import Q, QCombination
        except ImportError:
            from mongoengine.queryset.visitor import Q, QCombination
        from mongoengine.errors import ValidationError

        queries = map(lambda i: Q(**{i[0]: i[1]}), kwargs.items())
        query = QCombination(QCombination.AND, queries)
        try:
            return self.user_model.objects(query).first()
        except ValidationError:
            return None

    def find_role(self, role):
        return self.role_model.objects(name=role).first()

if __name__ == '__main__':
# Setup Flask-Security
user_datastore = MongoUserDatastore(db, User, Role)
security = Security(app, user_datastore)

# Create a user to test with
@app.before_first_request
def create_user():
    user_datastore.create_user(email='matt@nobien.net', password='password')

# Views
@app.route('/')
@login_required
def home():
    return render_template('index.html')

app.run()
